#define LIBTELNET_IMPLEMENTATION
#include "libtelnet.h"
#include <stdio.h>
#include <stdlib.h>
#include <signal.h>
#include <string.h>
#include <unistd.h>

/* Windows不支持termios.h，使用条件编译 */
#ifdef __unix__
#include <termios.h>
#include <poll.h>
#else
#include <conio.h>
#include <windows.h>
#endif

/* 全局会话变量，用于信号处理 */
telnet_session* g_session = NULL;
#ifdef __unix__
struct termios g_old_termios;
#endif
int g_running = 1;
static char input_buffer[1024];  // 行输入缓冲区
static size_t input_buffer_pos = 0;  // 行输入缓冲区位置

/* 恢复终端设置并清理 */
void cleanup(void) {
    if (g_session) {
        telnet_disconnect(g_session);
        telnet_destroy(g_session);
        g_session = NULL;
    }

#ifdef __unix__
    /* 恢复终端设置 */
    tcsetattr(STDIN_FILENO, TCSANOW, &g_old_termios);
#else
    /* 恢复Windows控制台模式 */
    HANDLE hStdin = GetStdHandle(STD_INPUT_HANDLE);
    DWORD mode;
    if (GetConsoleMode(hStdin, &mode)) {
        mode |= ENABLE_ECHO_INPUT | ENABLE_LINE_INPUT;
        SetConsoleMode(hStdin, mode);
    }
#endif
}

/* 信号处理函数 - 强制恢复终端模式 */
void signal_handler(int signum) {
    if (signum == SIGINT || signum == SIGTERM) {
        printf("\nReceived exit signal\n");
        cleanup();  // 确保清理终端模式
        g_running = 0;
        exit(0);
    }
}

/* 事件处理回调函数 */
void handle_event(telnet_session* session, const telnet_event* event) {
    switch (event->type) {
        case TELNET_EVENT_CONNECTED:
            printf("Connected to remote host\n");
            break;
            
        case TELNET_EVENT_DATA:
            /* 输出接收到的数据 */
            fwrite(event->event.data.data, 1, event->event.data.length, stdout);
            fflush(stdout);
            break;
            
        case TELNET_EVENT_DISCONNECTED:
            printf("\nConnection closed by remote host\n");
            g_running = 0;
            break;
            
        case TELNET_EVENT_ERROR:
            printf("\nError: %s\n", event->event.error.message);
            g_running = 0;
            break;
    }
}

/* 设置终端为合适的交互模式 */
void set_terminal_mode() {
#ifdef __unix__
    tcgetattr(STDIN_FILENO, &g_old_termios);
    struct termios new_termios = g_old_termios;
    
    /* 启用回显，禁用规范模式（不等待回车） */
    new_termios.c_lflag |= ECHO;       // 启用回显
    new_termios.c_lflag &= ~ICANON;    // 禁用规范模式
    new_termios.c_cc[VMIN] = 1;        // 至少读取1个字符
    new_termios.c_cc[VTIME] = 0;       // 不超时
    
    tcsetattr(STDIN_FILENO, TCSANOW, &new_termios);
#else
    /* Windows控制台设置 - 启用回显，禁用行缓冲 */
    HANDLE hStdin = GetStdHandle(STD_INPUT_HANDLE);
    DWORD mode;
    if (GetConsoleMode(hStdin, &mode)) {
        mode |= ENABLE_ECHO_INPUT;         // 启用回显
        mode &= ~ENABLE_LINE_INPUT;        // 禁用行输入（不等待回车）
        SetConsoleMode(hStdin, mode);
    }
#endif
}

/* 检查是否有用户输入 */
int has_user_input() {
#ifdef __unix__
    struct pollfd pfd;
    pfd.fd = STDIN_FILENO;
    pfd.events = POLLIN;
    return poll(&pfd, 1, 0) > 0;
#else
    return _kbhit();
#endif
}

/* 读取用户输入 - 修复行缓冲模式 */
int read_user_input() {
    if (input_buffer_pos >= sizeof(input_buffer) - 1) {
        // 缓冲区满，模拟回车处理
        input_buffer[input_buffer_pos] = '\n';
        input_buffer_pos++;
        int len = input_buffer_pos;
        input_buffer_pos = 0;
        return len;
    }

#ifdef __unix__
    char c;
    ssize_t n = read(STDIN_FILENO, &c, 1);
    if (n <= 0) return 0;
#else
    // Windows下读取字符
    char c = _getch();
#endif

    // 处理退格键
    if (c == '\b' || c == 127) {  // 127是ASCII DEL
        if (input_buffer_pos > 0) {
            input_buffer_pos--;
            // 从终端删除最后一个字符
            printf("\b \b");
            fflush(stdout);
        }
        return 0;
    }
    // 处理回车键
    else if (c == '\r' || c == '\n') {
        // 转换为换行符
        input_buffer[input_buffer_pos] = '\n';
        input_buffer_pos++;
        // 显示换行
        printf("\n");
        fflush(stdout);
        
        int len = input_buffer_pos;
        input_buffer_pos = 0;  // 重置缓冲区
        return len;
    }
    // 处理转义字符(^])
    else if (c == 0x1D) {
        return -1;  // 特殊标记表示转义字符
    }
    // 普通字符
    else {
        input_buffer[input_buffer_pos++] = c;
        // 回显字符（Unix下已由终端模式控制，Windows下需要手动回显）
#ifndef __unix__
        printf("%c", c);
        fflush(stdout);
#endif
        return 0;  // 未完成输入
    }
}

/* 显示帮助信息 */
void print_help(const char* program_name) {
    printf("Usage: %s [host [port]]\n", program_name);
    printf("Connect to a remote host via telnet\n");
    printf("Default port is 23 if not specified\n");
}

int main(int argc, char* argv[]) {
    const char* host = NULL;
    uint16_t port = 23;
    
    /* 解析命令行参数 */
    if (argc < 2) {
        print_help(argv[0]);
        return 1;
    } else if (argc >= 2) {
        host = argv[1];
        if (argc >= 3) {
            port = (uint16_t)atoi(argv[2]);
        }
    }
    
    /* 注册清理函数和信号处理 */
    atexit(cleanup);
    signal(SIGINT, signal_handler);
    signal(SIGTERM, signal_handler);
    
    /* 创建telnet会话 */
    g_session = telnet_create(4096);
    if (!g_session) {
        fprintf(stderr, "Failed to create telnet session\n");
        return 1;
    }
    
    /* 连接到远程主机 */
    printf("Trying to connect to %s:%d...\n", host, port);
    if (telnet_connect(g_session, host, port) != 0) {
        fprintf(stderr, "Failed to connect to %s:%d\n", host, port);
        return 1;
    }
    
    /* 设置终端为合适的交互模式 */
    set_terminal_mode();
    
    /* 主事件循环 */
    printf("Connected. Escape character is '^]'\n");
    while (g_running && g_session->state == TELNET_STATE_CONNECTED) {
        /* 处理telnet事件 */
        telnet_process(g_session, handle_event);
        
        /* 检查用户输入 */
        if (has_user_input()) {
            int result = read_user_input();
            
            if (result > 0) {
                /* 发送用户输入到服务器 */
                telnet_send(g_session, input_buffer, result);
            }
            // 处理转义字符(^])
            else if (result == -1) {
                printf("\nEscape character pressed. Enter 'quit' to exit, or 'send' to continue\n");
                
                /* 暂时恢复终端模式以读取命令 */
#ifdef __unix__
                tcsetattr(STDIN_FILENO, TCSANOW, &g_old_termios);
#else
                HANDLE hStdin = GetStdHandle(STD_INPUT_HANDLE);
                DWORD mode;
                if (GetConsoleMode(hStdin, &mode)) {
                    mode |= ENABLE_ECHO_INPUT | ENABLE_LINE_INPUT;
                    SetConsoleMode(hStdin, mode);
                }
#endif
                
                char cmd[32];
                if (fgets(cmd, sizeof(cmd), stdin)) {
                    if (strncmp(cmd, "quit", 4) == 0) {
                        g_running = 0;
                    }
                }
                
                /* 重新设置为交互模式 */
                set_terminal_mode();
                printf("Continuing connection...\n");
            }
        }
        
        /* 短暂休眠，减少CPU占用 */
#ifdef __unix__
        usleep(10000);
#else
        Sleep(10);
#endif
    }
    
    return 0;
}
